from typing import Dict, List
from .CodeParser import CodeParser

class CallGraphAnalyzer:
    """Analyzes function call relationships in parsed code."""
    
    def __init__(self, parser: CodeParser):
        self.functions = parser.functions
        self.entry_points = parser.entry_points
    
    def build_call_tree(self, direction: str = "down", max_depth: int = None) -> Dict:
        """
        Build a call tree for all functions in the program.
        
        Args:
            direction: "down" for functions called by each function, "up" for functions that call each function
            max_depth: Maximum depth of the tree (None for unlimited)
        
        Returns:
            A dictionary representing the call tree for the entire program
        """
        tree = {}
        visited = set()
        
        def build_subtree(node, current_depth, tree_dict):
            if max_depth is not None and current_depth > max_depth:
                return
            
            if node.full_name in visited:
                tree_dict[f"{node.full_name} (recursive)"] = {}
                return
            
            visited.add(node.full_name)
            subtree = {}
            
            if direction == "down":
                for callee in node.calls:
                    build_subtree(callee, current_depth + 1, subtree)
            elif direction == "up":
                for caller in node.called_by:
                    build_subtree(caller, current_depth + 1, subtree)
            
            visited.remove(node.full_name)
            tree_dict[node.full_name] = subtree
        
        # Build tree starting from each function
        for func_name, func in self.functions.items():
            visited = set()
            subtree = {}
            build_subtree(func, 1, subtree)
            tree.update(subtree)
            
        return tree

    def generate_call_chains(self, target_name: str = None, max_depth: int = None) -> List[List[str]]:
        """
        Generate call chains for all functions in the program.
        
        Args:
            target_name: Optional target function name to filter chains (None for all chains)
            max_depth: Maximum depth of chains (None for unlimited)
        
        Returns:
            A list of chains (each chain is a list of function names)
        """
        all_chains = []
        
        # If target is specified, validate it exists
        target = None
        if target_name:
            if target_name not in self.functions:
                print(f"Error: Function '{target_name}' not found")
                return []
            target = self.functions[target_name]
        
        # Process each function as a potential root
        for root_name, root in self.functions.items():
            chains = []
            
            def dfs(node, path, depth):
                if max_depth is not None and depth > max_depth:
                    return
                
                if target and node == target:
                    chains.append(path.copy())
                    return
                
                if not target and not node.calls:
                    chains.append(path.copy())
                    return
                
                for callee in node.calls:
                    if callee.full_name not in path:  # Avoid cycles
                        path.append(callee.full_name)
                        dfs(callee, path, depth + 1)
                        path.pop()
            
            dfs(root, [root.full_name], 1)
            all_chains.extend(chains)
        
        return all_chains

    def extract_computation_steps(self) -> Dict[str, str]:
        """
        Extract computation steps for all functions in the program.
        
        Returns:
            A dictionary mapping function names to their source code
        """
        steps = {}
        
        # Simply collect source code for all functions
        for func_name, func in self.functions.items():
            steps[func_name] = func.source_code
        
        return steps
